{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Estimating noise and stabilizing variance in calcium imaging data\n",
    "\n",
    "The purpose of this notebook is to show how to use the VST transform for estimating the noise profile of calcium imaging data and apply a generalized Anscombe transform that aims to transform the noise into white Gaussian."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    if __IPYTHON__:\n",
    "        get_ipython().run_line_magic('load_ext', 'autoreload')\n",
    "        get_ipython().run_line_magic('autoreload', '2')\n",
    "except NameError:\n",
    "    pass\n",
    "\n",
    "import logging\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import os\n",
    "import timeit\n",
    "\n",
    "logfile = None # Replace with a path if you want to log to a file\n",
    "logger = logging.getLogger('caiman')\n",
    "# Set to logging.INFO if you want much output, potentially much more output\n",
    "logger.setLevel(logging.WARNING)\n",
    "logfmt = logging.Formatter('%(relativeCreated)12d [%(filename)s:%(funcName)20s():%(lineno)s] [%(process)d] %(message)s')\n",
    "if logfile is not None:\n",
    "    handler = logging.FileHandler(logfile)\n",
    "else:\n",
    "    handler = logging.StreamHandler()\n",
    "handler.setFormatter(logfmt)\n",
    "logger.addHandler(handler)\n",
    "\n",
    "import caiman.external.houghvst.estimation as est\n",
    "from caiman.external.houghvst.gat import compute_gat, compute_inverse_gat\n",
    "import caiman as cm\n",
    "from caiman.paths import caiman_datadir\n",
    "# Set play_movies to False if you want to disable play of movies, e.g. for remote-hosted Jupyter environments\n",
    "play_movies = True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Below is a function that will compute and apply the transformation and its inverse. The underlying noise model is scaled Poisson plus Gaussian, i.e., the underlying fluorescence value $x$ is related to the observed value $y$ by the equation\n",
    "\n",
    "$$y = \\alpha*\\text{Poisson}(x) + \\varepsilon$$\n",
    "\n",
    "where $\\alpha$ is non-negative scalar, and $\\varepsilon \\sim \\mathcal{N}(\\mu,\\sigma^2)$ is distributed according to a Gaussian distribution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def main():\n",
    "    fnames = [os.path.join(caiman_datadir(), 'example_movies', 'demoMovie.tif')]\n",
    "\n",
    "    movie = cm.load(fnames)\n",
    "    movie = movie.astype(float)\n",
    "\n",
    "    # makes estimation numerically better:\n",
    "    movie -= movie.mean()\n",
    "\n",
    "    # use one every 200 frames\n",
    "    temporal_stride = 100\n",
    "    # use one every 8 patches (patches are 8x8 by default)\n",
    "    spatial_stride = 6\n",
    "\n",
    "    movie_train = movie[::temporal_stride]\n",
    "\n",
    "    t = timeit.default_timer()\n",
    "    estimation_res = est.estimate_vst_movie(movie_train, stride=spatial_stride)\n",
    "    print('\\tTime', timeit.default_timer() - t)\n",
    "\n",
    "    alpha = estimation_res.alpha\n",
    "    sigma_sq = estimation_res.sigma_sq\n",
    "\n",
    "    movie_gat = compute_gat(movie, sigma_sq, alpha=alpha)\n",
    "    # save movie_gat here\n",
    "    movie_gat_inv = compute_inverse_gat(movie_gat, sigma_sq, alpha=alpha,\n",
    "                                        method='asym')\n",
    "    # save movie_gat_inv here\n",
    "    return movie, movie_gat, movie_gat_inv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "movie, movie_gat, movie_gat_inv = main()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The transformed movie should have more uniform dynamic range (press `q` to exit):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if play_movies:\n",
    "    movie_gat.play(magnification=4, q_max=99.8)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The movie might appear more noisy but information is preserved as seen from the correlation image:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "CI = movie.local_correlations(swap_dim=False)\n",
    "CI_gat = movie_gat.local_correlations(swap_dim=False)\n",
    "plt.figure(figsize=(15,5))\n",
    "plt.subplot(1,2,1); plt.imshow(CI); plt.colorbar(); plt.title('Correlation Image (original)')\n",
    "plt.subplot(1,2,2); plt.imshow(CI_gat); plt.colorbar(); plt.title('Correlation Image (transformed)')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The noise estimates in space should also be more uniform"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sn = cm.source_extraction.cnmf.pre_processing.get_noise_fft(movie.transpose(1,2,0), noise_method='mean')[0]\n",
    "sn_gat = cm.source_extraction.cnmf.pre_processing.get_noise_fft(movie_gat.transpose(1,2,0), noise_method='mean')[0]\n",
    "# sn = np.std(movie.transpose(1,2,0), axis=-1)\n",
    "# sn_gat = np.std(movie_gat.transpose(1,2,0), axis=-1)\n",
    "plt.figure(figsize=(15,5))\n",
    "plt.subplot(1,2,1); plt.imshow(sn); plt.colorbar(); plt.title('Noise standard deviation (original)')\n",
    "plt.subplot(1,2,2); plt.imshow(sn_gat); plt.colorbar(); plt.title('Noise standard deviation (transformed)')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we apply the inverse transform we approximately get back the original movie (press `q` to exit):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "h = cm.concatenate([movie, movie_gat_inv], axis=2)\n",
    "#if play_movies:\n",
    "#    h.play(magnification=5, q_max=99.5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "caiman_pytorch2",
   "language": "python",
   "name": "caiman_pytorch2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
